#!/usr/bin/env python3.1
# -*- coding: utf-8 -*-

import sys
import re
import argparse
import math

class RunHead:
    """
    Container class for  header information of a single scintillate run
    """

    def __init__(self, txt):
        self.fullhead = []
        self.scint_thickness = 0
        self.source_distance = 0
        self.air_gap = 0
        self.shell_thickness = 0

        self.fullhead = txt
        for line in txt:
            if line.startswith("# scintThickness"):
                self.scint_thickness = float(re.match(
                        r"# scintThickness: ([-.0-9]+) mm", line
                    ).group(1))
            elif line.startswith("# sourceDistance"):
                self.source_distance = float(re.match(
                        r"# sourceDistance: ([-.0-9]+) mm", line).group(1)
                    )
            elif line.startswith("# airGap"):
                self.air_gap = float(re.match(
                        r"# airGap: ([-.0-9]+) mm", line).group(1)
                    )
            elif line.startswith("# shellThickness"):
                self.shell_thickness = float(re.match(
                        r"# shellThickness: ([-.0-9]+) mm", line).group(1)
                    )

        self.scinttop = self.air_gap + self.shell_thickness
        self.scintbottom = (self.scint_thickness +
                            self.air_gap + self.shell_thickness)


class Event:
    """
    Container class for a single gamma event (automatically populated from text
    by the constructor)
    """

    def __init__(self, txt):
        self.photons = []
        self.positions = []
        self.scintillationdepths = []
        self.hits = 0
        self.createdphotons = 0
        self._i = 0
        self._sorted = 0
        global createpos

        for line in txt:
            if line[0] == "#":
                m = re.match(("# scintillating: (?P<depenergy>[-.0-9]+) keV of "
                              "(?P<origenergy>[-.0-9]+) keV from parent ID "
                              "(?P<pid>[0-9]+) deposited at \((?P<x>[-.0-9]+) "
                              "mm, (?P<y>[-.0-9]+) mm, (?P<zcoord>[-.0-9]+) "
                              "mm\) producing (?P<num>[0-9]+) optical photons"),
                             line)
                if m != None:
                    self.scintillationdepths.append(float(m.group('zcoord')))
                    self.createdphotons += int(m.group('num'))
                m = re.match("# (?P<numhits>[0-9]+) Hits detected.", line)
                if m != None:
                    self.hits = int(m.group('numhits'))
                    self.photons = [0.0] * self.hits
                    self.positions = [[0.0, 0.0]] * self.hits
            else:
                parts = line[:-1].split("\t")
                #print(parts)
                self.photons[self._i] = float(parts[0])
                if createpos:
                    self.positions[self._i] = (float(parts[1]), float(parts[2]))
                self._i += 1
        #print(self.scintillationdepths)
        #print(self.positions)

    def _sort(self):
        if self._sorted == 0:
            self.photons.sort()
            self._sorted = 1

    def firstphoton(self):
        self._sort()
        if len(self.photons) > 0:
            return self.photons[0]
        else:
            return float('NaN')

    def quantile(self, q, k):
        """
        Method R-8 as described at
        http://en.wikipedia.org/wiki/Quantile#Estimating_the_quantiles_of_a_population
        q: used quantile (q-quantile, e.g. q=4 for quartiles)
        k: # of quantile (first, second, ...)
        """
        N = len(self.photons)
        if N < 1:
            return float('NaN')
        p = k/q
        h = (N + 1.0/3.0)*p + 1.0/3.0 - 1 # subtract 1 because of 0-index-start
        if h < 0:
            return self.photons[0]
        if h >= N - 1:
            return self.photons[-1]
        hfloor = math.floor(h)
        return (self.photons[hfloor] +
                (h - hfloor)*(self.photons[hfloor+1] - self.photons[hfloor]))


class TimeHistogram:
    """Container class for 1D time Histograms and related values"""

    def __init__(self, minimum, numbins, maximum):
        self.rightedges = []
        self.values = []
        self.overflow = 0
        self.underflow = 0
        self.minimum = 0
        self.maximum = 0
        self.numbins = 0
        self.numevents = 0       # events contributing to this histogram
        self.createdphotons = 0  # number of photons created in all
                                 #  scintillation processes of all contributing
                                 #  events
        self.detectedphotons = 0

        self.minimum = minimum
        self.numbins = numbins
        self.maximum = maximum
        self.rightedges = [(x + 1) * (self.maximum - self.minimum) /
                           numbins for x in range(numbins)]
        self.values = [0] * numbins

    def new_event(self, createdphotons, detectedphotons):
        self.numevents += 1
        self.createdphotons += createdphotons
        self.detectedphotons += detectedphotons

    def increment(self, bin):
        if bin >= self.maximum:
            self.overflow += 1
        elif bin < self.minimum:
            self.underflow += 1
        else:
            self.values[int((bin - self.minimum) * self.numbins //
                            (self.maximum-self.minimum))] += 1

    def printhist(self):
        print("# This histogram contains", self.numevents,
              "events which contain a total of", self.createdphotons,
              "created and", self.detectedphotons, "detected photons.", file=tf)
        print("# Underflow photons:", self.underflow, "Overflow photons:",
              self.overflow, file=tf)
        print("# Bin's right edge (ns)\tPhotons", file=tf)
        for i in range(len(self.rightedges)):
            print(self.rightedges[i], self.values[i], sep="\t", file=tf)


class PosHistogram2D:
    """Container class for 2D position Histograms and related values"""

    def __init__(
        self, xminimum, xnumbins, xmaximum, yminimum, ynumbins, ymaximum
    ):
        self.xcenters = []
        self.ycenters = []
        self.values = []
        self.xminimum = 0
        self.xmaximum = 0
        self.xnumbins = 0 # number of bins in one dimension
        self.yminimum = 0
        self.ymaximum = 0
        self.ynumbins = 0 # number of bins in one dimension
        self.numevents = 0
        self.createdphotons = 0
        self.detectedphotons = 0

        self._xrange = 0
        self._yrange = 0

        self.xminimum = xminimum
        self.xnumbins = xnumbins
        self.xmaximum = xmaximum
        self.yminimum = yminimum
        self.ynumbins = ynumbins
        self.ymaximum = ymaximum

        self._xrange = xmaximum - xminimum
        self._yrange = ymaximum - yminimum

        xbinwidth = (self._xrange) / xnumbins
        self.xcenters = [xminimum + xbinwidth / 2.0 + xbinwidth *
                         i for i in range(xnumbins)]
        ybinwidth = (self._yrange) / ynumbins
        self.ycenters = [yminimum + ybinwidth / 2.0 + ybinwidth *
                         i for i in range(ynumbins)]

        self.values = [[0]*ynumbins for x in range(xnumbins)]

    def new_event(self, createdphotons, detectedphotons):
        self.numevents += 1
        self.createdphotons += createdphotons
        self.detectedphotons += detectedphotons

    def increment(self, x, y):
        if (x >= self.xminimum and x < self.xmaximum and
            y > self.yminimum and y < self.ymaximum):
            xidx = int((x - self.xminimum) * self.xnumbins // self._xrange)
            yidx = int((y - self.yminimum) * self.ynumbins // self._yrange)
            self.values[xidx][yidx] += 1
        else:
            print("Position {x},{y} out of range!".format(x, y),
                  file=sys.stderr)

    def printhist(self):
        print("# This histogram contains", self.numevents,
              "events which contain a total of", self.createdphotons,
              "created and", self.detectedphotons, "detected photons.",
              file=pf)
        print("# x (mm)\ty (mm))\tPhotons", file=pf)
        for i in range(len(self.values)):
            for j in range(len(self.values[i])):
                print(self.xcenters[i], self.ycenters[j], self.values[i][j],
                      sep="\t", file=pf)


class Statistics:
    """
    Container class for figure of merit results. For each gamma event a row
    is added to this classes' lists
    """

    def __init__(self):
        self.firstphoton = []
        self.firstquartile = []
        self.createdphotons = []
        self.detectedphotons = []
        self.tenpercentile = []

    def add_event(self, ev):
        self.firstphoton.append(ev.firstphoton())
        self.firstquartile.append(ev.quantile(4, 1))
        self.createdphotons.append(ev.createdphotons)
        self.detectedphotons.append(len(ev.photons))
        self.tenpercentile.append(ev.quantile(10, 1))

    def printstat(self):
        print("# First Photon (ns)\tFirst Quartile (ns)\t10-Percentile "
              "(ns)\tDetected Photons\tCreated Photons", file=sf)
        for i in range(len(self.firstphoton)):
            print(self.firstphoton[i], self.firstquartile[i],
                  self.tenpercentile[i], self.detectedphotons[i],
                  self.createdphotons[i], sep="\t", file=sf)

#- relative varianz der photonenzahl
#- varianz der differenz zwischen erstem photon und erstem quartil
#- varianz der position des maximums
#- varianz des zeitpunkts des ersten photons.


def analyze(event):
    """
    Analyzes contents of an event block
    """

    global numevents
    numevents = numevents + 1
    if numevents % 100 == 0:
        print("Evaluated", numevents, "events.", file=sys.stderr)
    if (len(event.scintillationdepths) > 0 and
        len(event.photons) > 0):
        global runhead
        # compute average scintillation z coordinate
        scinttop = runhead[-1].scinttop
        scintbot = runhead[-1].scintbottom
        avgdepth = 0
        avgevts = 0
        for dep in event.scintillationdepths:
            if dep >= scinttop and dep <= scintbot:
                avgdepth += dep
                avgevts += 1
        if avgevts > 0:
            avgdepth = avgdepth / avgevts
        else:
            avgdepth = event.scintillationdepths[0]
        # select histogram
        if (avgdepth <
                runhead[-1].scint_thickness / 2 +
                runhead[-1].air_gap +
                runhead[-1].shell_thickness):
            global tophist
            hist = tophist
        else:
            global bottomhist
            hist = bottomhist
        # add photon to time histogram
        hist.new_event(event.createdphotons, event.hits)
        for ev in event.photons:
            hist.increment(ev)
        # add photon to position histogram
        global createpos
        if createpos:
            global incidencepos
            incidencepos.new_event(event.createdphotons, event.hits)
            for ev in event.positions:
                #print(ev)
                incidencepos.increment(ev[0], ev[1])
        # create statistics
        global createstat
        if createstat:
            global stat
            stat.add_event(event)


def read_block(file):
    """
    Reads a double-newline separated block from 'file' at once and starts
    content dependent analysis
    """

    result = []
    line = "\n"
    while line == "\n":
        line = file.readline()
    while 1:
        result.append(line)
        line = file.readline()
        if line == "":
            return 0
        if line == "\n":
            break
    # create and assign object
    if (result[0].startswith("# scintillating: ") or
            result[0] == "# 0 Hits detected.\n"):
        event = Event(result)
        analyze(event)
    elif result[0].startswith("# ====="):
        global runhead
        runhead.append(RunHead(result))

    return 1


###############################################################################
# directly executed part

# parse command line
parser = argparse.ArgumentParser(
    formatter_class=argparse.RawDescriptionHelpFormatter,
    description='Evaluates Scintillate! output files and creates histograms',
    epilog="""
The results can be plottet using e.g. the command:
    gnuplot -e " \\
    set term pdfcairo; set output \'out.pdf\'; \\
    set xlabel \'Time (ns)\'; set ylabel \'Photons\'; \\
    plot \'outfile-time.dat\' using 1:2 index 0 with steps title \'top\', \\
         \'hist-time.dat\' using 1:2 index 1 lt 2 with steps title \'bottom\'" \\
         && okular out.pdf &"""
    )
parser.add_argument('--min', default=0, type=float,
                    help='Lower limit of the histograms in ns (default: 0.0)')
parser.add_argument('--max', default=1, type=float,
                    help='Upper limit of the histograms in ns (default: 1.0)')
parser.add_argument('--bins', default=100, type=int,
                    help='Number of histogram bins (default: 100)')
parser.add_argument('--pattern', metavar='PAT', default='outfile',
                    help=('Pattern used for result file creation. This '
                          'pattern \is extended by "-times.dat" and '
                          '"-pos.dat" respectively (default: "outfile")'))
parser.add_argument('--stopafter', metavar='N', default=0, type=int,
                    help=('Evaluation stops after the provided number of data '
                          'sets was read. Use "0" to evaluate all sets from '
                          'all files provided'))
parser.add_argument('--statistics', action='store_true',
                    help=('Add this flag in order to produce statistical '
                          'output (i. e. compute variances and figures of '
                          'merit)'))
parser.add_argument('--positions', action='store_true',
                    help=('Add this flag in order to produce position '
                          'output (i. e. an incidence position histogram)'))
parser.add_argument('file', nargs='+',
                    help=('Name(s) of the input files to read. If more than '
                          'one filename is supplied the files are '
                          'concatenated as supplied. Use "-" for stdin'))
args = parser.parse_args()

if args.min < 0:
    sys.exit("Value supplied to --min must be larger than or equal to 0.0")
if args.max <= args.min:
    sys.exit("Value supplied to --max must be larger than the value supplied "
             "to --min")
if args.bins <= 0:
    sys.exit("Value supplied to --bins must be larger than 0")
if args.stopafter < 0:
    sys.exit("Value supplied to --stopafter must be larger than or equal to 0")

createstat = args.statistics
createpos = args.positions

# Create Histograms
tophist = TimeHistogram(args.min, args.bins, args.max)
bottomhist = TimeHistogram(args.min, args.bins, args.max)
stat = Statistics()

# Define size of 2D-Histogram for positions
incidencepos = PosHistogram2D(-52, 104, 52, -52, 104, 52)

# initialize global variables
runhead = []    # header of currently evaluated geant run
numevents = 0

# Open outfiles
tfname = args.pattern + "-time.dat"
tf = open(tfname, 'w')
print('Writing to file', tf.name)
if createpos:
    pfname = args.pattern + "-pos.dat"
    pf = open(pfname, 'w')
    print('Writing to file', pf.name)
if createstat:
    sfname = args.pattern + "-stat.dat"
    sf = open(sfname, 'w')
    print('Writing to file', sf.name)

# open and read infiles
for fn in args.file:
    print('Reading data file', fn, file=sys.stderr)
    if fn == "-":
        f = sys.stdin
    else:
        f = open(fn, 'r')

    while read_block(f):
        if args.stopafter > 0 and numevents >= args.stopafter:
            break

# create output
h = []
h = h + ["#\n# Read data from ", len(runhead), " file(s).\n"]
for filehead in runhead:
    h = h + filehead.fullhead
h[0:0] = ["# hist.py evaluation of a ", h[4][2:-1], " simulation run\n"]
th = h[:]
th.insert(3, "# time spectra\n")
print(*th, sep="", file=tf)
print("# analyzed a total of", numevents, "gammas", file=tf)
print("\n\n# First histogram containing interactions at the scintillator's "
      "top (source sided) half", file=tf)
tophist.printhist()
print("\n\n# Second histogram containing interactions at the scintillator's "
      "bottom (PMT sided) half", file=tf)
bottomhist.printhist()

if createpos:
    ph = h[:]
    ph.insert(3, "# position spectrum\n")
    print(*ph, sep="", file=pf)
    print("# analyzed a total of", numevents, "gammas", file=pf)
    print("\n\n# Hit position histogram of the cathode surface", file=pf)
    incidencepos.printhist()

if createstat:
    sh = h[:]
    sh.insert(3, "# statistical values\n")
    print(*sh, sep="", file=sf)
    print("# analyzed a total of", numevents, "gammas", file=sf)
    print("\n\n# Figures of merit", file=sf)
    stat.printstat()

